import fasttext
import chatbot.config as config


def build_classify_model(by_word=False):
    """
    :param by_word: 是否使用单个字作为特征
    :return:
    """
    path = config.classify_corpus_train_path if not by_word else config.classify_corpus_by_word_train_path
    # 修改这些参数，训练的model准确率不同，可以多测试找出一个最优的
    # wordNgrams,N-gram词袋模型，用连续的N个词语作为一个特征，不一定越多越好。
    model = fasttext.train_supervised(path, epoch=20, wordNgrams=2, minCount=1)
    model_path = config.classify_model_path if not by_word else config.classify_model_path_by_word
    model.save_model(model_path)


def get_classify_model(by_word=False):
    model_path = config.classify_model_path if not by_word else config.classify_model_path_by_word
    model = fasttext.load_model(model_path)
    return model


def eval(by_word=False):
    """
    模型评估，获取模型准确率
    :param by_word:
    :return:
    """
    model = get_classify_model(by_word)
    input = []
    target = []
    eval_data_path = config.classify_corpus_test_path if not by_word else config.classify_corpus_by_word_test_path
    for line in open(eval_data_path, encoding='utf-8').readlines():
        temp = line.split("__label__")
        if len(temp) < 2:
            continue
        input.append(temp[0].strip())
        target.append([temp[1].strip()])
    labels, acc_list = model.predict(input)
    sum = 0
    print(len(labels), len(target))
    for i, j in zip(labels, target):
        if i[0].replace("__label__", "") == j[0]:
            sum += 1
    acc = sum / len(labels)  # 平均准确率，相等的数据在全部数据中的占比
    return acc
